#include "AsyncClient.hpp"

#include "AsyncWebSocket.h"
#include "esp_timer.h"
#include "esp_log.h"
#include "mbedtls/sha1.h"
#include "mbedtls/base64.h"

#define TAG "AsyncWebSocket"
#define AWSC_PING_PAYLOAD   "ESPAsyncWebServer-PING"
#define AWSC_PING_PAYLOAD_LEN   22

const char * WS_STR_CONNECTION = "Connection";
const char * WS_STR_UPGRADE = "Upgrade";
const char * WS_STR_ORIGIN = "Origin";
const char * WS_STR_VERSION = "Sec-WebSocket-Version";
const char * WS_STR_KEY = "Sec-WebSocket-Key";
const char * WS_STR_PROTOCOL = "Sec-WebSocket-Protocol";
const char * WS_STR_ACCEPT = "Sec-WebSocket-Accept";
const char * WS_STR_UUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

size_t webSocketSendFrameWindow(AsyncClient* client)
{
    size_t space = client->get_send_buffer_size();
    return space < 9 ? 0 : space - 8;
}

/// @brief 构建并发送 WebSocket 帧
/// @param client
/// @param final 是否为最后一帧
/// @param opcode WebSocket消息操作码
/// @param mask 消息是否掩码保护
/// @param data 数据指针
/// @param len 计划发送的字节数
/// @return 发送数据的长度
size_t webSocketSendFrame(AsyncClient* client, bool final, uint8_t opcode, bool mask, uint8_t* data, size_t len)
{

    uint8_t headLen;
    if (len > 125) {
        headLen = mask ? 8 : 4;
    } else {
        headLen = mask ? 6 : 2;
    }

    size_t space = client->get_send_buffer_size();
    if (space >= headLen) {

        space -= headLen;
        if (len > space) {
            len = space;
        }

        auto* buf = new uint8_t[headLen];
        uint8_t* maskPtr = nullptr;
        if (buf == nullptr) {
            return 0;
        } else {
            buf[0] = final ? (0x80 | (0x0f & opcode)) : (0x0f & opcode);
            if (len < 126) {
                if (mask) {
                    buf[1] = 0x80 | len;
                    buf[2] = rand() & 0xff;
                    buf[3] = rand() & 0xff;
                    buf[4] = rand() & 0xff;
                    buf[5] = rand() & 0xff;
                    maskPtr = &buf[2];
                } else {
                    buf[1] = len;
                }
            } else {
                if (mask) {
                    buf[1] = 0x80 | 126;
                    buf[4] = rand() & 0xff;
                    buf[5] = rand() & 0xff;
                    buf[6] = rand() & 0xff;
                    buf[7] = rand() & 0xff;
                    maskPtr = &buf[4];
                } else {
                    buf[1] = 126;
                }
                buf[2] = (uint8_t)((len >> 8) & 0xff);
                buf[3] = (uint8_t)(len & 0xff);
            }
        }

        if (client->add((const char *)buf, headLen) != headLen) {
            delete[] buf;
            return 0;
        }

        if (mask) {
            for (size_t i = 0; i < len; i++) {
                data[i] ^= maskPtr[i & 0x03];
            }
        }
        delete[] buf;

        if (client->add((const char *)data, len) != len) {
            return 0;
        }
        if (!client->send()) {
            return 0;
        }
        return len;
    }

    return 0;
}

AsyncWebSocketMessageBuffer::AsyncWebSocketMessageBuffer()
    : m_data(nullptr)
    , m_len(0)
    , m_lock(false)
    , m_count(0)
{}

AsyncWebSocketMessageBuffer::AsyncWebSocketMessageBuffer(size_t size)
    : m_data(nullptr)
    , m_len(size)
    , m_lock(false)
    , m_count(0)
{
    if (size > 0) {
        m_data = new uint8_t[size + 1];
        if (m_data != nullptr) {
            m_data[size] = 0;
        }
    }
}

AsyncWebSocketMessageBuffer::AsyncWebSocketMessageBuffer(uint8_t* data, size_t size)
    : m_data(nullptr)
    , m_len(size)
    , m_lock(false)
    , m_count(0)
{
    if (data == nullptr) {
        return;
    } else {
        m_data = new uint8_t[size + 1];
        if (m_data != nullptr) {
            memcpy(m_data, data, m_len);
            m_data[size] = 0;
        }
    }
}

/// @brief 拷贝构造函数，只增加引用，不复制数据
/// @param copy
AsyncWebSocketMessageBuffer::AsyncWebSocketMessageBuffer(const AsyncWebSocketMessageBuffer &copy)
    : m_data(nullptr)
    , m_len(0)
    , m_lock(false)
    , m_count(0)
{
    m_len = copy.m_len;
    m_lock = copy.m_lock;
    m_count = 0;

    if (m_len > 0) {
        m_data = new uint8_t[m_len + 1];
        if (m_data != nullptr) {
            m_data[m_len] = 0;
            memcpy(m_data, copy.m_data, m_len);
        }
    }
}

/// @brief 移动构造函数
/// @param move
AsyncWebSocketMessageBuffer::AsyncWebSocketMessageBuffer(AsyncWebSocketMessageBuffer && move)
    : m_data(nullptr)
    , m_len(0)
    , m_lock(false)
    , m_count(0)
{
    m_len = move.m_len;
    m_lock = move.m_lock;
    m_count = move.m_count;

    if (move.m_data != nullptr) {
        m_data = move.m_data;
        move.m_data = nullptr;
    }
}

AsyncWebSocketMessageBuffer::~AsyncWebSocketMessageBuffer()
{
    if (m_data != nullptr) {
        delete[] m_data;
    }
}

bool AsyncWebSocketMessageBuffer::reserve(size_t size)
{
    if (m_data != nullptr) {
        delete[] m_data;
        m_data = nullptr;
    }

    m_data = new uint8_t[size + 1];

    if (m_data) {
        m_data[size] = 0;
        m_len = size;
        return true;
    } else {
        m_data = new uint8_t[m_len + 1];
        if (m_data != nullptr) {
            m_data[m_len] = 0;
        }
        return false;
    }
    return false;
}

/// @brief 基本消息构造函数
/// @param data 数据指针
/// @param len 数据长度
/// @param opcode 帧类型
/// @param mask 是否掩码
AsyncWebSocketBasicMessage::AsyncWebSocketBasicMessage(const char * data, size_t len, uint8_t opcode, bool mask)
    : m_len(len)
    , m_sent(0)
    , m_ack(0)
    , m_acked(0)
{
    m_opcode = opcode & 0x07;
    m_mask = mask;
    m_data = new uint8_t[len + 1];
    if (m_data == nullptr) {
        m_data = 0;
        m_status = WS_MSG_ERROR;
    } else {
        m_status = WS_MSG_SENDING;
        memcpy(m_data, data, len);
        m_data[len] = 0;
    }
}

/// @brief 基本消息构造函数
/// @param opcode 帧类型
/// @param mask 是否掩码
AsyncWebSocketBasicMessage::AsyncWebSocketBasicMessage(uint8_t opcode, bool mask)
    : m_len(0)
    , m_sent(0)
    , m_ack(0)
    , m_acked(0)
    , m_data(nullptr)
{
    m_opcode = opcode & 0x07;
    m_mask = mask;
}

AsyncWebSocketBasicMessage::~AsyncWebSocketBasicMessage()
{
    if (m_data != nullptr) {
        delete[] m_data;
    }
}

/// @brief 记录收到确认的字节数
/// @param len
/// @param time
void AsyncWebSocketBasicMessage::ack(size_t len, uint64_t time)
{
    m_acked += len;
    if (m_sent == m_len && m_acked == m_ack) {
        m_status = WS_MSG_SENT;
    }
}

/// @brief 发送WebSocket消息
/// @param client
/// @return 发送的字节数
size_t AsyncWebSocketBasicMessage::send(AsyncClient* client)
{
    if (m_status != WS_MSG_SENDING) {
        return 0;
    }
    if (m_acked < m_ack) {  // 流控
        return 0;
    }
    if (m_sent == m_len) {
        if (m_acked == m_ack) {
            m_status = WS_MSG_SENT;
        }
        return 0;
    }
    if (m_sent > m_len) {
        m_status = WS_MSG_ERROR;
        return 0;
    }

    size_t toSend = m_len - m_sent;
    size_t window = webSocketSendFrameWindow(client);

    if (window < toSend) {
        toSend = window;
    }

    m_sent += toSend;
    m_ack += toSend + ((toSend < 126) ? 2 : 4) + (m_mask * 4);

    bool final = (m_sent == m_len);
    uint8_t* dPtr = (uint8_t*)(m_data + (m_sent - toSend));
    uint8_t opCode = (toSend && m_sent == toSend) ? m_opcode : (uint8_t)WS_CONTINUATION;

    size_t sent = webSocketSendFrame(client, final, opCode, m_mask, dPtr, toSend);
    m_status = WS_MSG_SENDING;
    if (toSend && sent != toSend) {
        m_sent -= (toSend - sent);
        m_ack -= (toSend - sent);
    }

    return sent;
}

AsyncWebSocketMultiMessage::AsyncWebSocketMultiMessage(AsyncWebSocketMessageBuffer* buffer, uint8_t opcode, bool mask)
    : m_len(0)
    , m_sent(0)
    , m_ack(0)
    , m_acked(0)
    , m_WSbuffer(nullptr)
{
    m_opcode = opcode & 0x07;
    m_mask = mask;
    if (buffer) {
        m_WSbuffer = buffer;
        (*m_WSbuffer)++;
        m_data = buffer->get();
        m_len = buffer->length();
        m_status = WS_MSG_SENDING;
    } else {
        m_status = WS_MSG_ERROR;
    }
}

AsyncWebSocketMultiMessage::~AsyncWebSocketMultiMessage()
{
    if (m_WSbuffer) {
        (*m_WSbuffer)--;
    }
}

void AsyncWebSocketMultiMessage::ack(size_t len, uint64_t time)
{
    m_acked += len;
    if (m_sent >= m_len && m_acked >= m_ack) {
        m_status = WS_MSG_SENT;
    }
}

size_t AsyncWebSocketMultiMessage::send(AsyncClient* client)
{
    if (m_status != WS_MSG_SENDING) {
        return 0;
    }
    if (m_acked < m_ack) {  // 流控
        return 0;
    }
    if (m_sent == m_len) {
        m_status = WS_MSG_SENT;
        return 0;
    }
    if (m_sent > m_len) {
        m_status = WS_MSG_ERROR;
        return 0;
    }

    size_t toSend = m_len - m_sent;
    size_t window = webSocketSendFrameWindow(client);

    if (window < toSend) {
        toSend = window;
    }

    m_sent += toSend;
    m_ack += toSend + ((toSend < 126) ? 2 : 4) + (m_mask * 4);

    bool final = (m_sent == m_len);
    uint8_t* dPtr = (uint8_t*)(m_data + (m_sent - toSend));
    uint8_t opCode = (toSend && m_sent == toSend) ? m_opcode : (uint8_t)WS_CONTINUATION;

    size_t sent = webSocketSendFrame(client, final, opCode, m_mask, dPtr, toSend);
    m_status = WS_MSG_SENDING;
    if (toSend && sent != toSend) {
        m_sent -= (toSend - sent);
        m_ack -= (toSend - sent);
    }

    return sent;
}

AsyncWebSocket::AsyncWebSocket(const std::string &url)
    : m_buffers(LinkedList<AsyncWebSocketMessageBuffer*>([](AsyncWebSocketMessageBuffer * b)
{
    delete b;
}))
, m_url(url)
, m_clients(LinkedList<AsyncWebSocketClient*>([](AsyncWebSocketClient* c)
{
    delete c;
}))
, m_cNextID(1)
, m_enabled(true)
{
    m_eventHandler = nullptr;
}

AsyncWebSocket::~AsyncWebSocket()
{}

void AsyncWebSocket::_handleEvent(AsyncWebSocketClient* client, AwsEventType type, void* arg, uint8_t* data, size_t len)
{
    if (m_eventHandler != nullptr) {
        m_eventHandler(this, client, type, arg, data, len);
    }
}

void AsyncWebSocket::_addClient(AsyncWebSocketClient* client)
{
    m_clients.add(client);
}

void AsyncWebSocket::_handleDisconnect(AsyncWebSocketClient* client)
{
    m_clients.remove_first([ = ](AsyncWebSocketClient * c) {
        return c->id() == client->id();
    });
}

bool AsyncWebSocket::availableForWriteAll()
{
    for (const auto &c : m_clients) {
        if (c->queueIsFull()) {
            return false;
        }
    }
    return true;
}

bool AsyncWebSocket::availableForWrite(uint32_t id)
{
    for (const auto &c : m_clients) {
        if (c->queueIsFull() && (c->id() == id)) {
            return false;
        }
    }
    return true;
}

size_t AsyncWebSocket::count() const
{
    return m_clients.count_if([](AsyncWebSocketClient * c) {
        return c->status() == WS_CONNECTED;
    });
}

AsyncWebSocketClient* AsyncWebSocket::client(uint32_t id)
{
    for (const auto &c : m_clients) {
        if (c->id() == id && c->status() == WS_CONNECTED) {
            return c;
        }
    }
    return nullptr;
}

void AsyncWebSocket::close(uint32_t id, uint16_t code, const char* message)
{
    AsyncWebSocketClient* c = client(id);
    if (c != nullptr) {
        c->close(code, message);
    }
}

void AsyncWebSocket::closeAll(uint16_t code, const char* message)
{
    for (const auto &c : m_clients) {
        if (c->status() == WS_CONNECTED) {
            c->close(code, message);
        }
    }
}

void AsyncWebSocket::cleanupClients(uint16_t maxClients)
{
    if (count() > maxClients) {
        m_clients.front()->close();
    }
}

void AsyncWebSocket::ping(uint32_t id, uint8_t* data, size_t len)
{
    auto* c = client(id);
    if (c != nullptr) {
        c->ping(data, len);
    }
}

void AsyncWebSocket::pingAll(uint8_t* data, size_t len)
{
    for (const auto &c : m_clients) {
        if (c->status() == WS_CONNECTED) {
            c->ping(data, len);
        }
    }
}

void AsyncWebSocket::text(uint32_t id, const char* message)
{
    auto* c = client(id);
    if (c != nullptr) {
        c->text(message);
    }
}

void AsyncWebSocket::text(uint32_t id, const std::string &message)
{
    auto* c = client(id);
    if (c != nullptr) {
        c->text(message.c_str());
    }
}

void AsyncWebSocket::textAll(const char* message)
{
    auto* wsBuffer = makeBuffer((uint8_t*)message, strlen(message));
    textAll(wsBuffer);
}

void AsyncWebSocket::textAll(AsyncWebSocketMessageBuffer* buffer)
{
    if (buffer == nullptr) {
        return;
    }
    buffer->lock();
    for (const auto &c : m_clients) {
        if (c->status() == WS_CONNECTED) {
            c->text(buffer);
        }
    }
    buffer->unlock();
    _cleanBuffers();
}

void AsyncWebSocket::binary(uint32_t id, const char* message)
{
    auto* c = client(id);
    if (c != nullptr) {
        c->binary(message);
    }
}

void AsyncWebSocket::binary(uint32_t id, const std::string &message)
{
    auto* c = client(id);
    if (c != nullptr) {
        c->binary(message.c_str());
    }
}

void AsyncWebSocket::binaryAll(const char* message)
{
    auto* buffer = makeBuffer((uint8_t*)message, strlen(message));
    binaryAll(buffer);
}

void AsyncWebSocket::binaryAll(AsyncWebSocketMessageBuffer* buffer)
{
    if (buffer == nullptr) {
        return;
    }
    buffer->lock();
    for (const auto &c : m_clients) {
        if (c->status() == WS_CONNECTED) {
            c->binary(buffer);
        }
    }
    buffer->unlock();
    _cleanBuffers();
}

void AsyncWebSocket::message(uint32_t id, AsyncWebSocketMessage* message)
{
    auto* c = client(id);
    if (c) {
        c->message(message);
    }
}

void AsyncWebSocket::messageAll(AsyncWebSocketMessage* message)
{
    for (const auto &c : m_clients) {
        if (c->status() == WS_CONNECTED) {
            c->message(message);
        }
    }
    _cleanBuffers();
}

AsyncWebSocket::AsyncWebSocketClientLinkedList AsyncWebSocket::getClients() const
{
    return m_clients;
}

AsyncWebSocketMessageBuffer* AsyncWebSocket::makeBuffer(size_t size)
{
    auto* buffer = new AsyncWebSocketMessageBuffer(size);
    if (buffer) {
        AsyncWebLockGuard l(m_lock);
        m_buffers.add(buffer);
    }
    return buffer;
}

AsyncWebSocketMessageBuffer* AsyncWebSocket::makeBuffer(uint8_t* data, size_t size)
{
    auto* buffer = new AsyncWebSocketMessageBuffer(data, size);
    if (buffer) {
        AsyncWebLockGuard l(m_lock);
        m_buffers.add(buffer);
    }
    return buffer;
}

void AsyncWebSocket::_cleanBuffers()
{
    AsyncWebLockGuard l(m_lock);
    for (auto * c : m_buffers) {
        if (c && c->canDelete()) {
            m_buffers.remove(c);
        }
    }
}

void AsyncWebSocket::handleRequest(AsyncWebServerRequest* req)
{
    if (!req->hasHeader(std::string(WS_STR_VERSION)) || !req->hasHeader(std::string(WS_STR_KEY))) {
        req->send(400);
        return;
    }
    if ((m_username != "" && m_password != "") && !req->authenticate(m_username.c_str(), m_password.c_str())) {
        return req->requestAuthentication();
    }

    auto* version = req->getHeader(std::string(WS_STR_VERSION));
    if (std::stoi(version->value()) != 13) {
        auto* response = req->beginResponse(400);
        response->addHeader(std::string(WS_STR_VERSION), "13");
        req->send(response);
        return;
    }
    auto* key = req->getHeader(std::string(WS_STR_KEY));
    auto* response = new AsyncWebSocketResponse(key->value(), this);
    //.............
    if (req->hasHeader(std::string(WS_STR_PROTOCOL))) {
        auto* protocol = req->getHeader(WS_STR_PROTOCOL);
        response->addHeader(WS_STR_PROTOCOL, protocol->value());
    }
    req->send(response);
}

/// @brief 判断请求是否一个有效的WebSocket握手协议
/// @param req
/// @return 是则返回true，并在请求中添加感兴趣的头
bool AsyncWebSocket::canHandle(AsyncWebServerRequest* req)
{
    if (m_enabled == false) {
        return false;
    }
    if (req->method() != HTTP_GET || req->url() != m_url || !req->isExpectedRequestedConnType(RCT_WS)) {
        return false;
    }

    req->addInterestingHeader(WS_STR_CONNECTION);
    req->addInterestingHeader(WS_STR_UPGRADE);
    req->addInterestingHeader(WS_STR_ORIGIN);
    req->addInterestingHeader(WS_STR_VERSION);
    req->addInterestingHeader(WS_STR_KEY);
    req->addInterestingHeader(WS_STR_PROTOCOL);
    return true;
}

/// 控制帧
class AsyncWebSocketControl {
public:
    AsyncWebSocketControl(uint8_t opcode, uint8_t* data = nullptr, size_t len = 0, bool mask = false)
        : m_opcode(opcode)
        , m_len(len)
        , m_mask(len & mask)
        , m_finished(false)
    {
        if (data == nullptr) {
            m_len = 0;
        }
        if (m_len) {
            if (m_len > 125) {
                m_len = 125;    // 协议规定，控制帧最大长度为125且不能分片
            }
            m_data = new uint8_t[m_len];
            if (m_data == nullptr) {
                m_len = 0;
            } else {
                memcpy(m_data, data, len);
            }
        } else {
            m_data = nullptr;
        }
    }
    virtual ~AsyncWebSocketControl()
    {
        if (m_data != nullptr) {
            delete[] m_data;
        }
    }
    virtual bool finished() const
    {
        return m_finished;
    }
    uint8_t opcode()
    {
        return m_opcode;
    }
    uint8_t len()
    {
        return m_len + 2;
    }
    size_t send(AsyncClient* client)
    {
        m_finished = true;
        return webSocketSendFrame(client, true, m_opcode & 0x0f, m_mask, m_data, m_len);
    }
private:
    uint8_t     m_opcode;   // 控制帧的操作码
    uint8_t*    m_data;     // 有数载荷数据指针
    size_t      m_len;      // 有效载荷长度
    bool        m_mask;     // 是否启用掩码
    bool        m_finished; // 帧是否发送完成
};
///

/// @brief 异步WebSocket客户端
/// @param req 请求
/// @param server 服务器
AsyncWebSocketClient::AsyncWebSocketClient(AsyncWebServerRequest* req, AsyncWebSocket* server)
    : m_controlQueue(LinkedList<AsyncWebSocketControl*>([](AsyncWebSocketControl * c)
{
    delete c;
}))
, m_messageQueue(LinkedList<AsyncWebSocketMessage*>([](AsyncWebSocketMessage *m)
{
    delete m;
}))
, _tempObject(nullptr)
{
    m_client = req->client();
    m_server = server;
    m_clientID =  m_server->_getNextId();
    m_status = WS_CONNECTED;
    m_pstate = 0;
    m_lastMessageTime = esp_timer_get_time() / 1000;
    m_keepAlivePeriod = 0;
    m_client->set_rx_timeout_second(0);
    m_client->set_error_event_handler([](void* r, AsyncClient * c, err_t error) {
        ((AsyncWebSocketClient*)r)->_onError(error);
    },
    this
                                     );
    m_client->set_ack_event_handler([](void* r, AsyncClient * c, size_t len, uint64_t time) {
        ((AsyncWebSocketClient*)r)->_onAck(len, time);
    },
    this
                                   );
    m_client->set_disconnected_event_handler([](void* r, AsyncClient * c) {
        ((AsyncWebSocketClient*)r)->_onDisconnect();
        c->close();
    },
    this
                                            );
    m_client->set_timeout_event_handler([](void* r, AsyncClient * c, uint64_t time) {
        ((AsyncWebSocketClient*)r)->_onTimeout(time);
    },
    this
                                       );
    m_client->set_data_received_handler([](void* r, AsyncClient * c, void* buf, size_t len) {
        ((AsyncWebSocketClient*)r)->_onData(buf, len);
    },
    this
                                       );
    m_client->set_poll_event_handler([](void* r, AsyncClient * c) {
        ((AsyncWebSocketClient*)r)->_onPoll();
    },
    this
                                    );
    m_server->_addClient(this);
    m_server->_handleEvent(this, WS_EVT_CONNECT, req, nullptr, 0);
    delete req;
}

AsyncWebSocketClient::~AsyncWebSocketClient()
{
    m_messageQueue.free();
    m_controlQueue.free();
    m_server->_handleEvent(this, WS_EVT_DISCONNECT, nullptr, nullptr, 0);
}

ip_addr_t AsyncWebSocketClient::remoteIP()
{
    if (m_client == nullptr) {
        return IPADDR4_INIT(0);
    }
    return m_client->get_remote_IP();
}

uint16_t AsyncWebSocketClient::remotePort()
{
    if (m_client == nullptr) {
        return 0;
    }
    return m_client->get_remote_port();
}

/// @brief 按控制帧、数据帧的顺序发送消息
void AsyncWebSocketClient::m_runQueue()
{
    // 清理已n完成的消息帧
    while (!m_messageQueue.isEmpty() && m_messageQueue.front()->finished()) {
        m_messageQueue.remove(m_messageQueue.front());
    }

    if (!m_controlQueue.isEmpty() && (m_messageQueue.isEmpty() || m_messageQueue.front()->betweenFrames()) && webSocketSendFrameWindow(m_client) > (size_t)(m_controlQueue.front()->len() - 1)) {
        // 控制帧队列不为空、消息帧队列为空或处于帧间状态、发送窗口大于帧长
        m_controlQueue.front()->send(m_client);
    } else if (!m_messageQueue.isEmpty() && m_messageQueue.front()->betweenFrames() && webSocketSendFrameWindow(m_client)) {
        m_messageQueue.front()->send(m_client);
    }
}

void AsyncWebSocketClient::binary(const char* message)
{
    m_queueMessage(new AsyncWebSocketBasicMessage(message, strlen(message), WS_BINARY));
}

void AsyncWebSocketClient::binary(uint8_t* message, size_t len)
{
    m_queueMessage(new AsyncWebSocketBasicMessage((char*)message, len, WS_BINARY));
}

void AsyncWebSocketClient::binary(AsyncWebSocketMessageBuffer* buffer)
{
    m_queueMessage(new AsyncWebSocketMultiMessage(buffer, WS_BINARY));
}

void AsyncWebSocketClient::text(const char* message)
{
    m_queueMessage(new AsyncWebSocketBasicMessage(message, strlen(message)));
}

void AsyncWebSocketClient::text(uint8_t* message, size_t len)
{
    m_queueMessage(new AsyncWebSocketBasicMessage((char*)message, len));
}

void AsyncWebSocketClient::text(AsyncWebSocketMessageBuffer* buffer)
{
    m_queueMessage(new AsyncWebSocketMultiMessage(buffer));
}

/// @brief 向控制帧队列中添加一个PING帧
/// @param data
/// @param len
void AsyncWebSocketClient::ping(uint8_t* data, size_t len)
{
    if (m_status == WS_CONNECTED) {
        m_queueControl(new AsyncWebSocketControl(WS_PING, data, len));
    }
}

/// @brief 查看消息队列是否为满
/// @return
bool AsyncWebSocketClient::queueIsFull()
{
    if ((m_messageQueue.length() >= CONFIG_WS_MAX_QUEUE_MESSAGES) || (m_status != WS_CONNECTED)) {
        return true;
    }
    return false;
}

/// @brief 将消息添加到发送队列并尝试发送
/// @param dataMessage
void AsyncWebSocketClient::m_queueMessage(AsyncWebSocketMessage* dataMessage)
{
    if (dataMessage == nullptr) {
        return;
    }
    if (m_status != WS_CONNECTED) {
        delete dataMessage;
        return;
    }

    if (m_messageQueue.length() >= CONFIG_WS_MAX_QUEUE_MESSAGES) {
        delete dataMessage;
        ESP_LOGW(TAG, "Message queue is full.");
    } else {
        m_messageQueue.add(dataMessage);
    }
    if (m_client->get_send_buffer_size()) {
        m_runQueue();
    }
}

/// @brief 将控制帧加入到发送队列并尝试发送
/// @param controlMessage
void AsyncWebSocketClient::m_queueControl(AsyncWebSocketControl* controlMessage)
{
    if (controlMessage == nullptr) {
        return;
    }
    m_controlQueue.add(controlMessage);
    if (m_client->get_send_buffer_size()) {
        m_runQueue();
    }
}

/// @brief 构造一个关闭帧并发送出去
/// @param code 关闭代码
/// @param message 关闭消息
void AsyncWebSocketClient::close(uint16_t code, const char* message)
{
    if (m_status != WS_CONNECTED) {
        return;
    }

    if (code) {
        uint8_t packetLen = 2;
        if (message != nullptr) {
            size_t mlen = strlen(message);
            packetLen += (mlen > 123 ? 123 : mlen); // 控制帧最大125字节
        }

        auto* buf = new char[packetLen];
        if (buf != nullptr) {
            buf[0] = (uint8_t)(code >> 8);
            buf[1] = (uint8_t)(code & 0xff);
            if (message != nullptr) {
                memcpy(buf + 2, message, packetLen - 2);
            }
            m_queueControl(new AsyncWebSocketControl(WS_DISCONNECT, (uint8_t*)buf, packetLen));
            delete[] buf;
            return;
        }
    }
    m_queueControl(new AsyncWebSocketControl(WS_DISCONNECT));
}

/// @brief 处理客户端接收数据
/// @param pbuf
/// @param len
void AsyncWebSocketClient::_onData(void* pbuf, size_t len)
{
    m_lastMessageTime = esp_timer_get_time() / 1000;
    auto* data = (uint8_t*)pbuf;
    while (len > 0) {
        if (m_pstate == 0) { // TCP为流式协议，保证有序性但不保证一次性送达完整数据
            const uint8_t* frameData = data;
            m_pinfo.index = 0;
            m_pinfo.final = (frameData[0] & 0x80) != 0;     // FIN标志获取
            m_pinfo.opcode = frameData[0] & 0x0f;           // 获取当前帧操作码
            m_pinfo.masked = (frameData[1] & 0x80) != 0;    // 掩码标志获取
            m_pinfo.len = frameData[1] & 0x7f;              // 获取帧长度
            data += 2;
            len -= 2;

            // 帧长度调整
            if (m_pinfo.len == 126) {
                m_pinfo.len = ((uint16_t)(frameData[2] << 8)) | frameData[3];
                data += 2;
                len -= 2;
            } else if (m_pinfo.len == 127) {
                m_pinfo.len = ((uint64_t)(frameData[2] << 56)) |
                              ((uint64_t)(frameData[3] << 48)) |
                              ((uint64_t)(frameData[4] << 40)) |
                              ((uint64_t)(frameData[5] << 32)) |
                              ((uint64_t)(frameData[6] << 24)) |
                              ((uint64_t)(frameData[7] << 16)) |
                              ((uint64_t)(frameData[8] << 8)) |
                              frameData[9];
                data += 8;
                len -= 8;
            }

            if (m_pinfo.masked) {   // 获取掩码
                memcpy(m_pinfo.mask, data, 4);
                data += 4;
                len -= 4;
            }

            // 当前可处理长度调整
            const size_t dataLen = std::min((size_t)(m_pinfo.len - m_pinfo.index), len);
            const auto dataLast = data[dataLen];    // 保存最后一个数据

            // 数据解密
            if (m_pinfo.masked) {
                for (size_t i = 0; i < dataLen; i++) {
                    data[i] ^= m_pinfo.mask[(m_pinfo.index + i) & 0x03];
                }
            }

            if ((dataLen + m_pinfo.index) < m_pinfo.len) {          // 帧数据未处理完
                m_pstate = 1;
                if (m_pinfo.index == 0) {
                    if (m_pinfo.opcode) {   // 非继续帧
                        m_pinfo.message_opcode = m_pinfo.opcode;
                        m_pinfo.num = 0;
                    } else {                // 继续帧
                        m_pinfo.num += 1;
                    }
                    m_server->_handleEvent(this, WS_EVT_DATA, (void*)&m_pinfo, (uint8_t*)data, dataLen);
                    m_pinfo.index += dataLen;
                }
            } else if ((dataLen + m_pinfo.index) == m_pinfo.len) {  // 帧数据已处理完
                m_pstate = 0;
                if (m_pinfo.opcode == WS_DISCONNECT) {  // 关闭帧处理
                    if (dataLen) {
                        uint16_t reasonCode = (uint16_t)(data[0] << 8) + data[1];
                        char* reasonString = (char*)(data + 2);
                        if (reasonCode > 1001) { // 1000-4999为关闭码，大于1001为异常关闭
                            m_server->_handleEvent(this, WS_EVT_ERROR, (void*)&reasonCode, (uint8_t*)reasonString, strlen(reasonString));
                        }
                    }
                    if (m_status == WS_DISCONNECTING) {
                        m_status = WS_DISCONNECTED;     // 标记为断开
                        m_client->close(true);          // 关闭连接
                    } else {
                        m_status = WS_DISCONNECTING;    // 标记为正在断开
                        m_client->set_defer_ack(false); // 延迟确认
                        m_queueControl(new AsyncWebSocketControl(WS_DISCONNECT, data, dataLen));    // 队列化操作
                    }
                } else if (m_pinfo.opcode == WS_PING) {
                    m_queueControl(new AsyncWebSocketControl(WS_PONG, data, dataLen));
                } else if (m_pinfo.opcode == WS_PONG) {
                    if (dataLen != AWSC_PING_PAYLOAD_LEN || memcmp(AWSC_PING_PAYLOAD, data, AWSC_PING_PAYLOAD_LEN) != 0) {
                        m_server->_handleEvent(this, WS_EVT_PONG, nullptr, data, dataLen);
                    }
                } else if (m_pinfo.opcode < 8) {
                    // 继续帧/文本帧/二进制帧
                    m_server->_handleEvent(this, WS_EVT_DATA, (void*)&m_pinfo, data, dataLen);
                }
            } else {
                ESP_LOGE(TAG, "WebSocket Frame overflow!!!");
                break;
            }

            if (dataLen > 0) {
                data[dataLen] = dataLast;
            }
            data += dataLen;
            len -= dataLen;
        }
    }
}

/// @brief 轮询，有帧发送时则发送帧、无帧发送且超时时发送ping帧
void AsyncWebSocketClient::_onPoll()
{
    if (m_client->get_send_buffer_size()
            && (!m_controlQueue.isEmpty() || !m_messageQueue.isEmpty())) {
        m_runQueue();
    } else if (m_keepAlivePeriod > 0 && m_controlQueue.isEmpty()
               && m_messageQueue.isEmpty()
               && (esp_timer_get_time() / 1000 - m_lastMessageTime) >= m_keepAlivePeriod) {
        ping((uint8_t*)AWSC_PING_PAYLOAD, AWSC_PING_PAYLOAD_LEN);
    }
}

/// @brief
/// @param len
/// @param time
void AsyncWebSocketClient::_onAck(size_t len, uint64_t time)
{
    m_lastMessageTime = esp_timer_get_time() / 1000;
    if (!m_controlQueue.isEmpty()) {
        auto head = m_controlQueue.front();
        if (head->finished()) {
            len -= head->len();
            if (m_status == WS_DISCONNECTING && head->opcode() == WS_DISCONNECT) {
                m_controlQueue.remove(head);
                m_status = WS_DISCONNECTED;
                m_client->close(true);
                return;
            }
            m_controlQueue.remove(head);
        }
    }
    if (len && !m_messageQueue.isEmpty()) {
        m_messageQueue.front()->ack(len, time);
    }
    m_server->_cleanBuffers();
    m_runQueue();
}

/// @brief 处理超时的回调
/// @param time
void AsyncWebSocketClient::_onTimeout(uint64_t time)
{
    m_client->close(true);
}

/// @brief 处理断开连接的回调
void AsyncWebSocketClient::_onDisconnect()
{
    m_client = nullptr;
    m_server->_handleDisconnect(this);
}

AsyncWebSocketResponse::AsyncWebSocketResponse(const std::string &key, AsyncWebSocket* server)
    : m_server(server)
{
    m_code = 101;
    m_sendContentLength = false;

    auto* hash = new uint8_t[20];
    if (hash == nullptr) {
        m_state = RESPONSE_FAILED;
        return;
    }
    auto* buffer = new char[33];
    if (buffer == nullptr) {
        delete[] hash;
        m_state = RESPONSE_FAILED;
        return;
    }

    (std::string &)key += std::string(WS_STR_UUID);

    mbedtls_sha1_context    ctx;
    mbedtls_sha1_init(&ctx);
    mbedtls_sha1_starts(&ctx);
    mbedtls_sha1_update(&ctx, (const unsigned char*)key.c_str(), key.length());
    mbedtls_sha1_finish(&ctx, hash);
    mbedtls_sha1_free(&ctx);

    size_t encodeLength = 0;
    mbedtls_base64_encode((u_char*)buffer, 33, &encodeLength, hash, 20);
    addHeader(WS_STR_CONNECTION, WS_STR_UPGRADE);
    addHeader(WS_STR_UPGRADE, "websocket");
    addHeader(WS_STR_ACCEPT, buffer);
    delete[] buffer;
    delete[] hash;
}

void AsyncWebSocketResponse::_respond(AsyncWebServerRequest* req)
{
    if (m_state == RESPONSE_FAILED) {
        req->client()->close();
        return;
    }
    std::string out = _assembleHead(req->version());
    req->client()->write(out.c_str(), m_headLength);
    m_state = RESPONSE_WAIT_ACK;
}

size_t AsyncWebSocketResponse::_ack(AsyncWebServerRequest* req, size_t len, uint64_t time)
{
    if (len) {
        new AsyncWebSocketClient(req, m_server);
    }
    return 0;
}
